import pandas as pd
import numpy as np
import re
import statsmodels.formula.api as smf
from scipy import stats
from scipy.stats import sem

datadir = 'Folder with data'
whichtable = 'Table4'
#The regdatWeeksOut1 should be the design matrix used in the regressions for fatalities 1 week out.
regdat = pd.read_csv(datadir+'regdatWeeksOut1_'+whichtable+'.csv',low_memory=False,dtype={'countyFIPS':'str'})
#Match paper's normalization of the GrowthFatRate.
regdat['GrowthFatRate']=regdat['GrowthFatRate']/1000

regformula = 'GrowthFatRate ~ 1 + HHMF + HHSS + TEMPcoldMF + TEMPcoldSS + TEMP + PerCapIncome +\
Age65plus + Age85plus + Asian + Black + Hispanic + NativeAmer + Other + Diabetes +\
Obese + Smoke + HousingDensity + PopDensity + NurseTotPop + w + tMinusTrig +\
TotDeathsPerCap+dGatherMax100Start + dbarCrestC + dRisk2Closed + dResidentMaskMandStart +\
dGatherMax100plusStart + dGatherMax10Start + dNursingMustAcceptStart + dNursingHomeStart +\
dGymsZeroStart + dElectiveStart + dParksStart + dSpasZeroStart + dRisk1Closed + dbarCrestOpen +\
dResidentMaskRecStart + dBusOpenRevStart + dRisk3Closed + dRisk4Closed + dStayHomeStart + dEmergencyStart + dEmployeeMaskStart'



resultAll = smf.ols(regformula, regdat).fit()

outorder = {
'dStayHomeStart':'Stay at Home',
'dEmergencyStart':'State of Emergency',
'dNursingMustAcceptStart':'Nursing Home Accept Pos.',
'dNursingHomeStart':'No Nursing Home Visit',
'dEmployeeMaskStart':'Employee Mask',
'dResidentMaskRecStart':'Masks Recommended',
'dResidentMaskMandStart':'Mandatory Masks',
'dParksStart':'Parks Closed',
'dElectiveStart':'No Elective Procedures',
'dbarCrestC':'Restaurants and Bars Closed', 
'dbarCrestOpen':'Bars Closed/Restaurants Open',
'dGymsZeroStart':'Gyms Closed',  
'dSpasZeroStart':'Spas Closed',
'dGatherMax10Start':'Gatherings limited to 10',
'dGatherMax100Start':'No gatherings over 100',
'dGatherMax100plusStart':'Gathering limit over 100',
'dRisk1Closed':'Risk Level 1 to 4 Closed',
'dBusOpenRevStart':'Bus Openings Reversed'}

residByOrder = pd.DataFrame(index=[outorder[x] for x in outorder],columns=['Mean1','SE1','p-value1','Sig.*1','Mean2','SE2','p-value2','Sig.*2'])
regformula = re.sub(' +',' ',regformula)

for i in outorder:
      regresid = re.sub('\+ '+i+' \+','+',regformula)
      result = smf.ols(regresid, regdat).fit()
      zz = result.resid
      zz = pd.DataFrame(zz,columns=['residual'])
      zz[[x for x in regdat if x.startswith('w') or x.startswith('d') and 'End' not in x]] = regdat[[x for x in regdat if x.startswith('w') or x.startswith('d') and 'End' not in x]]
      for j in [1,2]:
             residWeekX = zz.loc[(zz['w'+i[1:]]>=4*j/7) & (zz['w'+i[1:]]<((4*j/7)+1)),['residual']]
             outMean = residWeekX.mean()[0]
             pval = stats.ttest_1samp(residWeekX,0)[1][0]
             residByOrder.loc[outorder[i],'Mean'+str(j)] = outMean.round(3)
             residByOrder.loc[outorder[i],'SE'+str(j)] = sem(residWeekX)[0].round(3)
             residByOrder.loc[outorder[i],'p-value'+str(j)] = pval.round(3)
             if pval<0.01:
                     residByOrder.loc[outorder[i],'Sig.*'+str(j)] = '***'
             elif (pval>=0.01) & (pval<0.05):
                     residByOrder.loc[outorder[i],'Sig.*'+str(j)] = '**'
             elif pval<0.1:
                     residByOrder.loc[outorder[i],'Sig.*'+str(j)] = '*'
             else:
                     residByOrder.loc[outorder[i],'Sig.*'+str(j)] = ''

residByOrder.drop(columns=['p-value1','p-value2'],inplace=True)
print(residByOrder)




